package com.nacol.demo.gateway.loadbalance;

import cn.hutool.core.util.StrUtil;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Random;
import java.util.stream.Collectors;

/**
 * @Author Nacol(姚秋实)
 * @Date 2023/12/11
 * @Title 自定义负载均衡
 * @Description 目的：该类仅演示负载均衡使用 Instance.MetaData 作为负载均衡条件
 */
public class CustomLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private final String serviceId;
    private final DiscoveryClient discoveryClient;
    private final Random random = new Random();

    public CustomLoadBalancer(String serviceId, LoadBalancerClientFactory loadBalancerClientFactory) {
        this.serviceId = serviceId;
        this.discoveryClient = loadBalancerClientFactory.getLazyProvider(serviceId, DiscoveryClient.class).getIfAvailable();
    }

    /**
     * @Author Nacol(姚秋实)
     * @Date 2023/12/11
     * @Title 自定义负载均衡
     * @Description 仅作为DEMO演示使用metadata控制负载均衡
     * ApiPost传送门：https://console-docs.apipost.cn/preview/39d3272b69136477/e0d238a2b63a570f
     */
    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        // 从Nacos获取服务实例列表
        List<ServiceInstance> instances = discoveryClient.getInstances(serviceId);

        /**
         * 过滤只对【customFlag=green】的节点进行负载均衡
         * 此处可扩展自己的需求
         */
        instances = instances.stream()
                .filter(instance -> StrUtil.equals(instance.getMetadata().get("customFlag"), "green"))
                .collect(Collectors.toList());

        if (instances.isEmpty()) {
            return Mono.empty();
        }

        // 随机选择一个实例
        ServiceInstance instance = instances.get(random.nextInt(instances.size()));
        return Mono.just(new DefaultResponse(instance));
    }

}